import traceback

import arrow
from pymongo import MongoClient, UpdateOne, InsertOne

from conf.settings import CLOUDMONITOR_MONGODB
from utility.utils import backup_warp
from utility.log import Logger

logger = Logger().get_logger()

collection_tuple = ('packages', 'charges')
# db = MongoClient(host='localhost', port=27017)['cloudmonitor']
db = MongoClient(CLOUDMONITOR_MONGODB)['cloudmonitor']

arrow_now = arrow.utcnow()
now_datetime = arrow_now.datetime


@backup_warp(collection_tuple, logger, db, arrow_now)
def main():
    plan = db.packages.find_one({'key': 'scanv_basic'})
    plan_id = plan['_id']

    plan_total = {}
    for k, v in plan['limit'].items():
        if v.get('total', 0) != 0:
            plan_total[k] = v

    uid_map = {item.get('uid'): item for item in db.balances.find({'planId': plan_id, 'planExpireTime': {'$gt': arrow_now.datetime}})}

    charge_update_list = []
    balance_update_list = []
    for uid, balance in uid_map.items():
        balance_item = balance['balance']
        balance_update_dict = {}
        plan_charge_time = balance['planChargeTime']
        plan_expire_time = balance['planExpireTime']
        for k, v in plan_total.items():
            # if "ManualTask" in k:
            #     v['limitType'] = 'period'

            query = {'chargeType': 'plan', 'uid': uid, 'funcType': k, 'isDelete': False}

            if not (charge := db.charges.find_one(query)):
                time = 0
                if v.get('limitType') == 'period':
                    time = arrow_now.replace(day=1, hour=2, minute=0, second=0, microsecond=0).shift(months=1, hours=-8, minutes=-5).int_timestamp
                insert_dict = {
                    "limitType": v['limitType'],
                    "used": 0,
                    "total": v.get('total', 0),
                    "number": 0,
                    "funcType": k,
                    "uid": uid,
                    "chargeType": "plan",
                    "time": time,
                    "createTime": plan_charge_time,
                    "updateTime": now_datetime,
                    "taskType": v.get('taskType', ''),
                    "expireTime": plan_expire_time,
                    "isDelete": False,
                    "packId": plan_id
                }
                charge_update_list.append(InsertOne(insert_dict))
            elif charge['total'] != v['total']:
                update_dict = {'total': v.get('total', 0), 'updateTime': now_datetime}
                if not charge.get('packId'):
                    update_dict['packId'] = plan_id

                charge_update_list.append(UpdateOne(query, {'$set': update_dict}, upsert=True))

            if not balance_item.get(k):
                balance_update_dict.setdefault(f'balance', {})[k] = {
                    'used': 0,
                    'total': v.get('total', 0),
                    'time': 0,
                    'limitType': v.get('limitType', ''),
                    'taskType': v.get('taskType', ''),
                    'number': 0
                }

        if balance_update_dict:
            balance_update_list.append(UpdateOne({'_id': balance['_id']}, {'$set': balance_update_dict}))

    if balance_update_list:
        db.balance.bulk_write(balance_update_list)
    if charge_update_list:
        db.charges.bulk_write(charge_update_list)


if __name__ == '__main__':
    main()
